import torch
import torch.nn as nn
import os
import random
from tqdm import tqdm
from torch.utils.data import DataLoader
from dataset import MyDataset
from diffusion import UNet
from pred import prediction

if __name__ == '__main__':
    img_rootPath = r'../JPEGImages'
    saveModel_dir = r'checkpoints'
    verification = r'verification'
    load_model = r''
    epoch = 300
    batch_size = 8
    lr = 1e-4
    alpha = 4e-4
    
        
    dataset = MyDataset(img_rootPath)
    train_dl = DataLoader(dataset, shuffle=True, batch_size=batch_size, drop_last=True)


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = UNet()
    model.to(device)

    optim = torch.optim.Adam(model.parameters(), lr)
    loss_fn = nn.MSELoss()
    
    for e in range(epoch):
        model.train()
        with tqdm(train_dl) as train_dl_buffer:
            
            main_total_loss = 0
            for batch, (xt_img, mask, timesteps, target_epsilon) in enumerate(train_dl_buffer):
                xt_img = xt_img.to(device)
                mask = mask.to(device)
                timesteps = timesteps.to(device)
                target_epsilon = target_epsilon.to(device)

                pred_epsilon = model(xt_img, mask, timesteps)
                loss = loss_fn(pred_epsilon, target_epsilon)

                optim.zero_grad()
                loss.backward()
                optim.step()

                main_total_loss += loss.item()
                batch += 1
                train_dl_buffer.set_description("Epoch {}: loss {:.8f}".format(e+1, main_total_loss/batch))

        # 保存模型 
        path = os.path.join(saveModel_dir, 'checkpoint.pth')
        torch.save(model.state_dict(), path)
        print("保存模型===》{}".format(path))

        for i in range(10):
            img_path = dataset.img_path_arr[random.randint(0, dataset.__len__()-1)]
            save_path = os.path.join(verification, f'{i}.jpg')
            prediction(model, img_path, save_path, device)